#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "aes.h" 


int hex_char_to_int(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
}

int hex_str_to_bytes(const char *hex, uint8_t *bytes, size_t max_len) {
    size_t len = strlen(hex);
    if (len % 2 != 0 || (len / 2) > max_len) {
        fprintf(stderr, "Error: Invalid hex string length (%zu) or exceeds max_len (%zu).\n", len, max_len);
        return -1;
    }
    for (size_t i = 0; i < len; i += 2) {
        int high = hex_char_to_int(hex[i]);
        int low = hex_char_to_int(hex[i + 1]);
        if (high == -1 || low == -1) {
            fprintf(stderr, "Error: Invalid character in hex string: %c%c\n", hex[i], hex[i+1]);
            return -1;
        }
        bytes[i / 2] = (uint8_t)((high << 4) | low);
    }
    return (int)(len / 2);
}

void bytes_to_hex_str(const uint8_t *bytes, size_t len, char *hex_str) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex_str + (i * 2), "%02x", bytes[i]);
    }
    hex_str[len * 2] = '\0';
}

// --- 核心測試 ---
int test_aes(const char *test_name, const char *pt_hex, const char *key_hex, const char *ct_hex_expected) {
    printf("[ RUN      ] %s\n", test_name);

    uint8_t key_bytes[32];
    uint8_t plaintext[AES_BLOCK_SIZE];
    uint8_t expected_ciphertext[AES_BLOCK_SIZE];
    uint8_t buffer[AES_BLOCK_SIZE];
    char actual_ct_hex[AES_BLOCK_SIZE * 2 + 1];

    int key_len_bytes = hex_str_to_bytes(key_hex, key_bytes, sizeof(key_bytes));
    if (key_len_bytes <= 0) {
        return 0;
    }

    if (hex_str_to_bytes(pt_hex, plaintext, sizeof(plaintext)) != AES_BLOCK_SIZE) {
        return 0;
    }

    if (hex_str_to_bytes(ct_hex_expected, expected_ciphertext, sizeof(expected_ciphertext)) != AES_BLOCK_SIZE) {
        return 0;
    }

    AES_KeySize key_type_enum;
    if (key_len_bytes == 16) key_type_enum = AES_KEY_128;
    else if (key_len_bytes == 24) key_type_enum = AES_KEY_192;
    else if (key_len_bytes == 32) key_type_enum = AES_KEY_256;
    else {
        fprintf(stderr, "[  FAILED  ] %s - Invalid key size: %d bytes\n", test_name, key_len_bytes);
        return 0;
    }

    const int Nr = AES_Nr(key_type_enum);
    const int total_words = AES_Nb * (Nr + 1);

    uint32_t w_encrypt[total_words];
    uint32_t w_decrypt[total_words];

    if (AES_KeyExpansion(key_bytes, w_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "[  FAILED  ] %s - Key expansion for encrypt failed\n", test_name);
        return 0;
    }

    if (AES_KeyExpansion(key_bytes, w_decrypt, key_type_enum, AES_KEY_SCHEDULE_DECRYPT) != 0) {
        fprintf(stderr, "[  FAILED  ] %s - Key expansion for decrypt failed\n", test_name);
        return 0;
    }

    memcpy(buffer, plaintext, sizeof(plaintext));
    AES_Encrypt(buffer, w_encrypt, key_type_enum);

    if (memcmp(buffer, expected_ciphertext, sizeof(expected_ciphertext)) != 0) {
        bytes_to_hex_str(buffer, sizeof(buffer), actual_ct_hex);
        printf("[  FAILED  ] %s - Encryption mismatch\n", test_name);
        printf("  Expected: %s\n", ct_hex_expected);
        printf("  Actual  : %s\n", actual_ct_hex);
        return 0;
    }

    AES_Decrypt(buffer, w_decrypt, key_type_enum);

    if (memcmp(buffer, plaintext, sizeof(plaintext)) != 0) {
        printf("[  FAILED  ] %s - Decryption mismatch (Plaintext not recovered)\n", test_name);
        bytes_to_hex_str(buffer, sizeof(buffer), actual_ct_hex);
        printf("  Original Plaintext: %s\n", pt_hex);
        printf("  Decrypted Result  : %s\n", actual_ct_hex);
        return 0;
    }

    printf("[       OK ] %s\n", test_name);
    return 1;
}

// --- 測試函式 ---
void benchmark_custom_aes() {
    // 變數宣告區 
    const AES_KeySize key_type_enum = AES_KEY_256;
    const size_t data_size = 16 * 1024 * 1024; // 16 MB
    const int iterations = 10;
    const int Nr = AES_Nr(key_type_enum);
    const int total_words = AES_Nb * (Nr + 1);

    // 將 VLA 宣告移動到函式開頭
    uint32_t w_encrypt[total_words];
    uint32_t w_decrypt[total_words];

    uint8_t *key_bytes = NULL;
    uint8_t *plaintext = NULL;
    uint8_t *ciphertext = NULL;
    uint8_t *decryptedtext = NULL;
    int error_occurred = 0;

    // 將所有計時和效能變數的宣告提前到函式開頭
    clock_t start_encrypt = 0, end_encrypt = 0;
    clock_t start_decrypt = 0, end_decrypt = 0;
    double total_data_mb = 0;
    double encrypt_duration = 0;
    double encrypt_speed_mbps = 0;
    double decrypt_duration = 0;
    double decrypt_speed_mbps = 0;

    // --- 資源分配與初始化 ---
    key_bytes = (uint8_t*)malloc(key_type_enum * 4);
    plaintext = (uint8_t*)malloc(data_size);
    ciphertext = (uint8_t*)malloc(data_size);
    decryptedtext = (uint8_t*)malloc(data_size);

    if (!key_bytes || !plaintext || !ciphertext || !decryptedtext) {
        fprintf(stderr, "Memory allocation failed for benchmark.\n");
        error_occurred = 1;
        
        goto cleanup; 
    }

    srand(time(NULL));
    for (int i = 0; i < (key_type_enum * 4); ++i) key_bytes[i] = rand() % 256;
    for (size_t i = 0; i < data_size; ++i) plaintext[i] = rand() % 256;

    if (AES_KeyExpansion(key_bytes, w_encrypt, key_type_enum, AES_KEY_SCHEDULE_ENCRYPT) != 0) {
        fprintf(stderr, "Benchmark key expansion for encrypt failed\n");
        error_occurred = 1;
        goto cleanup;
    }
    if (AES_KeyExpansion(key_bytes, w_decrypt, key_type_enum, AES_KEY_SCHEDULE_DECRYPT) != 0) {
        fprintf(stderr, "Benchmark key expansion for decrypt failed\n");
        error_occurred = 1;
        goto cleanup;
    }

    // --- 加密測試 ---
    printf("\nStarting encryption benchmark (AES-%d)...\n", key_type_enum * 32);
    start_encrypt = clock();
    for (int i = 0; i < iterations; ++i) {
        uint8_t *temp_buffer = (uint8_t*)malloc(data_size);
        if (!temp_buffer) {
            fprintf(stderr, "Memory allocation failed during encryption benchmark.\n");
            error_occurred = 1;
            goto cleanup_loop;
        }
        memcpy(temp_buffer, plaintext, data_size);
        for (size_t offset = 0; offset < data_size; offset += AES_BLOCK_SIZE) {
            AES_Encrypt(temp_buffer + offset, w_encrypt, key_type_enum);
        }
        if (i == 0) memcpy(ciphertext, temp_buffer, data_size);
        free(temp_buffer);
    }
cleanup_loop:
    end_encrypt = clock();
    if (error_occurred) goto cleanup;
    printf("Encryption finished.\n");

    // --- 解密測試 ---
    printf("Starting decryption benchmark (AES-%d)...\n", key_type_enum * 32);
    start_decrypt = clock();
    for (int i = 0; i < iterations; ++i) {
        uint8_t *temp_buffer = (uint8_t*)malloc(data_size);
        if (!temp_buffer) {
            fprintf(stderr, "Memory allocation failed during decryption benchmark.\n");
            error_occurred = 1;
            goto cleanup_loop_decrypt;
        }
        memcpy(temp_buffer, ciphertext, data_size);
        for (size_t offset = 0; offset < data_size; offset += AES_BLOCK_SIZE) {
            AES_Decrypt(temp_buffer + offset, w_decrypt, key_type_enum);
        }
        if (i == 0) memcpy(decryptedtext, temp_buffer, data_size);
        free(temp_buffer);
    }
cleanup_loop_decrypt:
    end_decrypt = clock();
    if (error_occurred) goto cleanup;
    printf("Decryption finished.\n");

if (memcmp(plaintext, decryptedtext, data_size) != 0) {
    printf("\nBenchmark Verification: FAILED!\n");
    
    int error_count = 0;
    for (size_t offset = 0; offset < data_size; offset += AES_BLOCK_SIZE) {
        if (memcmp(plaintext + offset, 
                   decryptedtext + offset, 
                   AES_BLOCK_SIZE) != 0) {
            error_count++;
            printf("發現損壞區塊 @ 偏移:%zu\n", offset);
            
            // 打印原始明文和解密結果的完整狀態
            uint32_t plain_state[4], decrypt_state[4];
            memcpy(plain_state, plaintext + offset, AES_BLOCK_SIZE);
            memcpy(decrypt_state, decryptedtext + offset, AES_BLOCK_SIZE);
            
            printf("  原始狀態: ");
            AES_PrintState(plaintext + offset);
            printf("  解密狀態: ");
            AES_PrintState(decryptedtext + offset);
            
            // 打印差異字節
            const uint8_t *p_plain = plaintext + offset;
            const uint8_t *p_decrypt = decryptedtext + offset;
            for (int i = 0; i < AES_BLOCK_SIZE; i++) {
                if (p_plain[i] != p_decrypt[i]) {
                    printf("    字節 %02d: 0x%02x → 0x%02x (差 %02x)\n", 
                           i, p_plain[i], p_decrypt[i],
                           p_plain[i] ^ p_decrypt[i]);
                }
            }
            
            // 限制打印的錯誤區塊數量
            if (error_count >= 5) {
                printf("... 已發現 %d 個錯誤區塊（僅顯示前5個）\n", error_count);
                break;
            }
        }
    }
} else {
    printf("\nBenchmark Verification: SUCCESS!\n");
}
        // --- 調試碼塊結束 ---
    

    total_data_mb = (double)data_size * iterations / (1024 * 1024);
    encrypt_duration = (double)(end_encrypt - start_encrypt) / CLOCKS_PER_SEC;
    encrypt_speed_mbps = (encrypt_duration > 0) ? (total_data_mb / encrypt_duration) : 0;
    decrypt_duration = (double)(end_decrypt - start_decrypt) / CLOCKS_PER_SEC;
    decrypt_speed_mbps = (decrypt_duration > 0) ? (total_data_mb / decrypt_duration) : 0;

    printf("\n--- Custom AES Benchmark Result ---\n");
    printf("--------------------------------------------------------\n");
    printf("Mode           : AES-%d-ECB (for benchmark only)\n", key_type_enum * 32);
    printf("Implementation : Pure C/Software (T-Table)\n");
    printf("Total Data     : %.2f MB\n", total_data_mb);
    printf("--------------------------------------------------------\n");
    printf("Encryption Time : %.2f s\n", encrypt_duration);
    printf("Encryption Speed: %.2f MB/s\n", encrypt_speed_mbps);
    printf("\n");
    printf("Decryption Time : %.2f s\n", decrypt_duration);
    printf("Decryption Speed: %.2f MB/s\n", decrypt_speed_mbps);
    printf("--------------------------------------------------------\n");

cleanup:
    free(key_bytes);
    free(plaintext);
    free(ciphertext);
    free(decryptedtext);
}


// --- 主函式 ---
int main() {
    int all_tests_passed = 1;

    printf("========================================\n");
    printf("     Running NIST AES Test Vectors      \n");
    printf("========================================\n");

    all_tests_passed &= test_aes("NIST F.1.1 AES-128", "6bc1bee22e409f96e93d7e117393172a", "2b7e151628aed2a6abf7158809cf4f3c", "3ad77bb40d7a3660a89ecaf32466ef97");
    all_tests_passed &= test_aes("NIST F.1.3 AES-192", "6bc1bee22e409f96e93d7e117393172a", "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b", "bd334f1d6e45f25ff712a214571fa5cc");
    all_tests_passed &= test_aes("NIST F.1.5 AES-256", "6bc1bee22e409f96e93d7e117393172a", "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4", "f3eed1bdb5d2a03c064b5a7e3db181f8");
    
    printf("========================================\n");

    if (all_tests_passed) {
        printf("All NIST validation tests passed!\n");
        benchmark_custom_aes();
    } else {
        fprintf(stderr, "One or more NIST validation tests failed. Skipping benchmark.\n");
        return 1;
    }

    return 0;
}
